# -*- coding: utf-8 -*-
from datetime import datetime

import pandas as pd
import json

from elasticsearch import Elasticsearch

from conf.variable import ES_HOSTS
from icd_kg.icd_kg_creator import IcdKgCreator

es = Elasticsearch(hosts=ES_HOSTS)

icd_kg_creator_app = IcdKgCreator()

if __name__ == '__main__':
    with open("../icd_api/data/icd11_2020_09_backup.json", "r", encoding="utf-8") as f:
        data = json.load(f)

    coln = list(data[0].keys())

    size = len(data)

    DF = pd.DataFrame(data, columns=coln)
    DF = DF.set_index("uid")

    count = 0
    for index, row in DF.iterrows():
        count += 1

        if index is None:
            continue

        print(f'Adding index to ES, count is {count}')

        index_term = list(row['index_term'])
        name = row['name']
        index_term.append(name)
        index_term = list(set(index_term))
        class_kind = row['class_kind']

        doc = {
            'uid': index,
            'index_term': index_term,
            'class_kind': class_kind,
            'timestamp': datetime.now(),
        }

        # res = es.index(index="cdss-icd11-index-term", id=index, body=doc)

    count = 0
    for index, row in DF.iterrows():
        count += 1
        if index is None:
            index = ''
        print(f'Adding node, count is {count}')
        # icd_kg_creator_app.add_node(index,
        #                             child=row['child'], parent=row['parent'],
        #                             properties={
        #                                 'name': row['name'],
        #                                 'code': row['code'],
        #                                 'class_kind': row['class_kind']
        #                             })

    count = 0
    for index, row in DF.iterrows():
        count += 1
        if index is None:
            index = ''
        if len(row['child']) > 0:
            print(f'Adding child, count is {count}')
            # icd_kg_creator_app.add_child(index)
        if len(row['parent']) > 0:
            # 此处需要手动设置一下根节点，将 parent 为 [] 替换为 ['']
            print(f'Adding parent, count is {count}')
            # icd_kg_creator_app.add_parent(index)
